import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

# args = parser.parse_args()
# if args.dataset == 'cifar100':
#     num_classes = 100
# elif args.dataset == 'cifar10':
#     num_classes = 10
# elif args.dataset == 'tinyimagenet':
#     num_classes = 200
# else: #iNat18
#     num_classes = 8142

def ib_loss(input_values, ib):
    """Computes the focal loss"""
    loss = input_values * ib
    return loss.mean()

class IBLoss(nn.Module):
    def __init__(self, weight=None, alpha=10000., num_classes=7):
        super(IBLoss, self).__init__()
        assert alpha > 0
        self.alpha = alpha
        self.epsilon = 0.001
        self.weight = weight
        self.num_classes = num_classes

    def forward(self, input, target, features):
        target = target.long()
        grads = torch.sum(torch.abs(F.softmax(input, dim=1) - F.one_hot(target, self.num_classes)),1) # N * 1
        ib = grads*features.reshape(-1)
        ib = self.alpha / (ib + self.epsilon)
        return ib_loss(F.cross_entropy(input, target, reduction='none', weight=self.weight), ib)

def ib_focal_loss(input_values, ib, gamma):
    """Computes the ib focal loss"""
    p = torch.exp(-input_values)
    ib = ib.unsqueeze(1)
    loss = (1 - p) ** gamma * input_values * ib
    return loss.mean()

class IB_FocalLoss(nn.Module):
    def __init__(self, weight=None, alpha=10000., gamma=0., num_classes=7):
        super(IB_FocalLoss, self).__init__()
        assert alpha > 0
        self.alpha = alpha
        self.epsilon = 0.001
        self.weight = weight
        self.gamma = gamma
        self.num_classes = num_classes

    def forward(self, input, target, features):
        # grads = torch.sum(torch.abs(F.softmax(input, dim=1) - F.one_hot(target, self.num_classes)),1) # N * 1       
        # 使用 sigmoid 函数将 input 转换为概率
        probs = torch.sigmoid(input)

        # 计算每个类别上的梯度
        grads = torch.abs(probs - target)  # N * num_classes

        # 对每个样本的梯度求和
        grads = torch.sum(grads, dim=1)  # N * 1
        ib = grads * (features.reshape(-1))
        ib = self.alpha / (ib + self.epsilon)
        # print('input: ',input.shape)
        # print('target: ',target.shape)
        # print('ib: ',ib.shape)
        return ib_focal_loss(F.binary_cross_entropy_with_logits(input, target, reduction='none', weight=self.weight), ib, self.gamma)

def focal_loss(input_values, gamma):
    """Computes the focal loss"""
    p = torch.exp(-input_values)
    #loss = (1 - p) ** gamma * input_values
    loss = (1- p) ** gamma * input_values * 10
    return loss.mean()

class FocalLoss(nn.Module):
    def __init__(self, weight=None, gamma=0.):
        super(FocalLoss, self).__init__()
        assert gamma >= 0
        self.gamma = gamma
        self.weight = weight

    def forward(self, input, target):
        return focal_loss(F.cross_entropy(input, target, reduction='none', weight=self.weight), self.gamma)

class LDAMLoss(nn.Module):
    
    def __init__(self, cls_num_list, max_m=0.5, weight=None, s=30):
        super(LDAMLoss, self).__init__()
        m_list = 1.0 / np.sqrt(np.sqrt(cls_num_list))
        m_list = m_list * (max_m / np.max(m_list))
        m_list = torch.cuda.FloatTensor(m_list)
        self.m_list = m_list
        assert s > 0
        self.s = s
        self.weight = weight

    def forward(self, x, target):
        index = torch.zeros_like(x, dtype=torch.uint8)
        index.scatter_(1, target.data.view(-1, 1), 1)
        
        index_float = index.type(torch.cuda.FloatTensor)
        batch_m = torch.matmul(self.m_list[None, :], index_float.transpose(0,1))
        batch_m = batch_m.view((-1, 1))
        x_m = x - batch_m
    
        output = torch.where(index, x_m, x)
        return F.cross_entropy(self.s*output, target, weight=self.weight)

class MultiFocalLoss(nn.Module):
    def __init__(self, gamma=2, alpha=None, reduction='mean'):
        super(MultiFocalLoss, self).__init__()
        self.gamma = gamma
        self.alpha = alpha
        self.reduction = reduction

    def forward(self, inputs, targets):
        """
        计算多标签 Focal Loss

        Args:
            inputs:  模型输出的 logits，shape 为 (batch_size, num_classes)
            targets: 真实标签，shape 为 (batch_size, num_classes)，值应为 0 或 1
            gamma:   Focal Loss 的聚焦参数 (gamma >= 0)
            alpha:   类别权重，可以是一个 list 或 tensor，shape 为 (num_classes)，或者 None 表示不使用类别权重
            reduction: Loss 的 reduction 方式，'none' | 'mean' | 'sum'，'mean' 表示对 batch size 求平均，'sum' 表示求和，'none' 表示不进行 reduction

        Returns:
            loss:  计算得到的 Focal Loss，是一个标量 (如果 reduction='mean' 或 'sum') 或 shape 为 (batch_size, num_classes) 的 tensor (如果 reduction='none')
        """
        inputs = torch.sigmoid(inputs) #  确保 inputs 经过 sigmoid 激活，得到概率值

        bce_loss = F.binary_cross_entropy(inputs, targets, reduction='none') # 先计算 BCE Loss

        pt = torch.where(targets == 1, inputs, 1 - inputs) #  pt = p if y=1, pt = 1-p if y=0
        focal_weight = (1 - pt) ** self.gamma # 计算调制系数 (1-pt)^gamma

        loss = focal_weight * bce_loss # 应用调制系数到 BCE Loss

        if self.alpha is not None: # 应用类别权重 alpha
            # alpha_factor = torch.tensor(self.alpha).to(inputs.device).unsqueeze(0) #  转换为 tensor 并 unsqueeze 适应 batch_size
            alpha_factor = self.alpha.clone().detach().to(inputs.device).unsqueeze(0)
            loss = alpha_factor * loss

        if self.reduction == 'mean':
            loss = loss.mean()
        elif self.reduction == 'sum':
            loss = loss.sum()
        elif self.reduction == 'none':
            pass #  保持原样，不进行 reduction

        return loss